#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import sys, os

null_fd = os.open(os.devnull, os.O_WRONLY)
old_stderr = os.dup(2)  
os.dup2(null_fd, 2)  

import cv2

class TraditionalTrafficLightDetector:
    """
    基于传统计算机视觉的红绿灯检测节点
    使用颜色分割和形状检测来识别红绿灯
    """
    
    def __init__(self):
        
        # 调试模式参数
        self.debug_mode = True
        self.binary_show = True
        
        # 颜色范围 (HSV) - 优化后的范围，减少光晕影响
        self.color_ranges = {
            'red': [
                # 红色范围1: 0-10度
                (np.array([0, 120, 120]), np.array([10, 255, 255])),
                # 红色范围2: 160-180度
                (np.array([160, 120, 120]), np.array([180, 255, 255]))
            ],
            'green': [
                # 绿色范围: 提高饱和度和亮度下限，避免光晕
                (np.array([50, 100, 100]), np.array([80, 255, 255]))
            ],
            'yellow': [
                # 黄色范围: 添加黄灯检测
                (np.array([20, 120, 120]), np.array([30, 255, 255]))
            ]
        }
        
    
    def image_callback(self, cv_image):
        """
        图像回调函数
        """
        # 执行红绿灯检测
        detection_result, total_area, annotated_image = self.detect_traffic_light(cv_image)
        
        # 根据检测结果发布信号
        if detection_result is not None:
            # 如果看到红灯，发送False；否则发送True
            go_signal = False if detection_result == 'red' else True
            
            # 中文状态映射
            status_text = "可通行" if go_signal else "停止"
            color_text = "绿灯" if detection_result == 'green' else "红灯"
            signal_text = "True" if go_signal else "False"
            
            print(f"检测到: {color_text}, 状态: {status_text}, 发送信号: {signal_text}, 总面积: {total_area:.1f}")
        else:
            # 什么都没看到，发送True（默认可通行）
            go_signal = True
            print(f"未检测到红绿灯，发送信号: True")
        
        # 显示检测结果图像 
        if self.debug_mode:
            # 确保窗口名称一致
            window_name = 'Traffic Light Detection'
            
            # 调整图像大小便于显示（如果图像太大）
            height, width = annotated_image.shape[:2]
            if width > 800:
                scale = 800.0 / width
                new_width = int(width * scale)
                new_height = int(height * scale)
                annotated_image = cv2.resize(annotated_image, (new_width, new_height))
                print(f"[DEBUG] 调整显示图像尺寸: {width}x{height} -> {new_width}x{new_height}")
            
            # 使用OpenCV显示图像
            cv2.imshow(window_name, annotated_image)
            
            # 设置窗口属性
            cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE)
            
            # 确保图像被刷新显示
            cv2.waitKey(1)  # 非阻塞等待，允许图像更新
            print(f"[DEBUG] 图像已显示在窗口: {window_name}")
    
    def detect_traffic_light(self, image):
        """
        主要的红绿灯检测函数 - 基于轮廓面积的简化检测
        返回: (颜色名称, 轮廓面积, 标注图像)
        """
        height, width = image.shape[:2]
        annotated_image = image.copy()
        
        print(f"开始检测红绿灯，图像尺寸: {width}x{height}")
        
        # 光晕消除预处理
        processed_image = self.remove_halo_effect(image)
        
        # 转换到HSV颜色空间
        hsv = cv2.cvtColor(processed_image, cv2.COLOR_BGR2HSV)
        
        # 应用高斯模糊减少噪声
        hsv_blurred = cv2.GaussianBlur(hsv, (9, 9), 2)
        
        # 只检测红灯和绿灯
        detection_results = {}
        color_masks = {}
        
        # 检测红灯和绿灯
        for color_name in ['red', 'green']:
            if color_name not in self.color_ranges:
                continue
                
            ranges = self.color_ranges[color_name]
            print(f"\n检测 {color_name} 颜色...")
            
            contours, total_area, mask = self.detect_contours_for_color(hsv_blurred, ranges, color_name)
            
            # 存储掩码用于调试显示
            if self.binary_show:
                color_masks[color_name] = mask
            
            # 调试信息
            mask_pixels = cv2.countNonZero(mask)
            mask_percentage = (mask_pixels / (width * height)) * 100
            print(f"{color_name}: 掩码像素数量: {mask_pixels} ({mask_percentage:.2f}%)")
            
            if contours is not None:
                print(f"{color_name}: 检测到 {len(contours)} 个轮廓, 总面积: {total_area:.1f}")
                detection_results[color_name] = {
                    'contours': contours,
                    'total_area': total_area,
                    'mask': mask
                }
            else:
                print(f"{color_name}: 未检测到轮廓")
                detection_results[color_name] = {
                    'contours': None,
                    'total_area': 0,
                    'mask': mask
                }
        
        # 显示颜色掩码
        if self.binary_show:
            self.display_color_masks(color_masks)
        
        # 根据简化逻辑判断结果
        final_detection = self.determine_traffic_light_state(detection_results)
        
        # 绘制检测结果
        if final_detection and self.debug_mode:
            self.draw_contour_detection(annotated_image, detection_results[final_detection], final_detection)
        
        # 返回结果
        if final_detection is None:
            print("\n最终结果: 未检测到任何红绿灯")
            return None, 0, annotated_image
        else:
            total_area = detection_results[final_detection]['total_area']
            print(f"\n最终结果: 检测到 {final_detection} 灯，总面积: {total_area:.1f}")
        
        return final_detection, total_area, annotated_image
    
    def detect_contours_for_color(self, hsv_image, color_ranges, color_name):
        """
        使用轮廓检测方法检测红绿灯，返回轮廓列表和总面积
        """
        # 创建颜色掩码
        mask = np.zeros(hsv_image.shape[:2], dtype=np.uint8)
        for lower, upper in color_ranges:
            color_mask = cv2.inRange(hsv_image, lower, upper)
            mask = cv2.bitwise_or(mask, color_mask)
        
        # 形态学操作去除噪声和填充空洞
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        
        # 使用轮廓检测方法
        contours = self.detect_contours_for_traffic_lights(mask, color_name)
        
        if contours:
            # 计算所有轮廓的总面积
            total_area = self.calculate_total_contour_area(contours)
            return contours, total_area, mask
        
        return None, 0, mask
    
    def display_color_masks(self, color_masks):
        """
        显示三种颜色的二值化掩码图像
        """
        # 显示红色掩码
        if 'red' in color_masks:
            cv2.imshow('Red Mask', color_masks['red'])
        
        # 显示黄色掩码
        if 'yellow' in color_masks:
            cv2.imshow('Yellow Mask', color_masks['yellow'])
        
        # 显示绿色掩码
        if 'green' in color_masks:
            cv2.imshow('Green Mask', color_masks['green'])
        
        cv2.waitKey(1)  # 非阻塞等待，允许图像更新

    def calculate_confidence(self, mask, circles):
        """
        计算检测置信度
        """
        if len(circles) == 0:
            return 0
        
        # 对于每个检测到的圆，计算其在掩码中的像素密度
        confidences = []
        for (x, y, r) in circles:
            # 创建圆形ROI
            circle_mask = np.zeros(mask.shape, dtype=np.uint8)
            cv2.circle(circle_mask, (x, y), r, 255, -1)
            
            # 计算圆形区域内的像素密度
            roi_mask = cv2.bitwise_and(mask, circle_mask)
            pixel_count = cv2.countNonZero(roi_mask)
            total_pixels = np.pi * r * r
            
            if total_pixels > 0:
                density = pixel_count / total_pixels
                # 将密度转换为百分比置信度
                confidence = min(100, density * 150)  # 调整系数
                confidences.append(confidence)
        
        return max(confidences) if confidences else 0
    
    def draw_detection(self, image, x, y, r, color_name, confidence):
        """
        在图像上绘制检测结果
        """
        print(f"[DEBUG] draw_detection调用: color={color_name}, center=({x},{y}), radius={r}, confidence={confidence:.1f}%")
        
        # 颜色映射
        color_map = {
            'red': (0, 0, 255),
            'yellow': (0, 255, 255),
            'green': (0, 255, 0)
        }
        
        color = color_map.get(color_name, (255, 255, 255))
        print(f"[DEBUG] 使用绘制颜色: {color} for {color_name}")
        
        # 确保坐标为整数
        x, y, r = int(x), int(y), int(r)
        
        # 绘制圆形检测框（粗线）
        cv2.circle(image, (x, y), r, color, 4)
        print(f"[DEBUG] 绘制外圆: center=({x},{y}), radius={r}")
        
        # 绘制圆心（实心圆）
        cv2.circle(image, (x, y), 3, color, -1)
        print(f"[DEBUG] 绘制圆心: center=({x},{y})")
        
        # 添加标签
        label = f"{color_name}: {confidence:.1f}%"
        print(f"[DEBUG] 添加标签: {label}")
        
        label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2)[0]
        
        # 确保标签在图像范围内
        img_height, img_width = image.shape[:2]
        label_x = max(5, min(x - label_size[0] // 2, img_width - label_size[0] - 5))
        label_y = max(35, y - r - 15)
        
        print(f"[DEBUG] 标签位置: ({label_x}, {label_y}), 图像尺寸: {img_width}x{img_height}")
        
        # 绘制标签背景（黑色背景）
        cv2.rectangle(image, 
                     (label_x - 5, label_y - 28), 
                     (label_x + label_size[0] + 5, label_y + 5), 
                     (0, 0, 0), -1)
        
        # 绘制标签文字（白色文字）
        cv2.putText(image, label, (label_x, label_y - 5), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
        
        print(f"[DEBUG] draw_detection完成")
    

    def run(self):
        """
        运行红绿灯检测，调用image_callback函数
        """
        # 图像路径 - 请修改为您的实际图像路径
        image_path = "/home/windiff/self-control-homework/src/virtual_camera/test_data/images/3890c3235168436a4dde7c89036efa22.jpg"
        
        # 读取图像
        image = cv2.imread(image_path)
        
        if image is None:
            print(f"错误: 无法读取图像文件 {image_path}")
            print("请检查文件路径是否正确")
            return
        
        print(f"成功读取图像: {image_path}")
        print(f"图像尺寸: {image.shape}")
        
        # 显示原始图像
        cv2.imshow('Original Image', image)
        
        # 直接调用image_callback函数处理图像
        print("调用image_callback函数...")
        self.image_callback(image)
        
        print("按任意键退出...")
        cv2.waitKey(0)  # 等待按键
        cv2.destroyAllWindows()  # 关闭所有窗口

    def remove_halo_effect(self, image):
        """
        消除LED灯光晕效果的预处理方法
        """
        print("[预处理] 开始消除光晕效果...")
        
        # 方法1: 伽马校正 - 压缩高亮区域
        gamma_corrected = self.gamma_correction(image, gamma=0.6)
        
        # 方法2: 限制对比度自适应直方图均衡化(CLAHE)
        lab = cv2.cvtColor(gamma_corrected, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        
        # 对亮度通道应用CLAHE
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        l = clahe.apply(l)
        
        # 重新合并通道
        processed = cv2.merge([l, a, b])
        processed = cv2.cvtColor(processed, cv2.COLOR_LAB2BGR)
        
        # 方法3: 形态学操作去除小的光晕区域
        processed = self.morphological_halo_removal(processed)
        
        # 方法4: 基于饱和度的光晕抑制
        processed = self.saturation_based_halo_suppression(processed)
        
        print("[预处理] 光晕消除完成")
        
        # 可选：显示预处理效果对比
        if self.debug_mode:
            cv2.imshow('Original vs Processed', np.hstack([
                cv2.resize(image, (400, 300)),
                cv2.resize(processed, (400, 300))
            ]))
        
        return processed
    
    def gamma_correction(self, image, gamma=0.6):
        """
        伽马校正 - 压缩高亮区域，增强暗部细节
        gamma < 1: 压缩高亮
        gamma > 1: 增强对比度
        """
        # 构建查找表
        inv_gamma = 1.0 / gamma
        table = np.array([((i / 255.0) ** inv_gamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
        
        # 应用伽马校正
        return cv2.LUT(image, table)
    
    def morphological_halo_removal(self, image):
        """
        使用形态学操作去除光晕
        """
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 检测过亮区域
        _, bright_mask = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
        
        # 形态学开运算去除小的亮点
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
        bright_mask = cv2.morphologyEx(bright_mask, cv2.MORPH_OPEN, kernel, iterations=2)
        
        # 对过亮区域进行处理
        result = image.copy()
        bright_pixels = bright_mask > 0
        
        # 对过亮区域应用更强的模糊
        blurred = cv2.GaussianBlur(image, (15, 15), 5)
        result[bright_pixels] = blurred[bright_pixels]
        
        return result
    
    def saturation_based_halo_suppression(self, image):
        """
        基于饱和度的光晕抑制
        光晕通常具有高亮度但低饱和度的特征
        """
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)
        
        # 检测高亮度但低饱和度的区域（典型光晕特征）
        halo_mask = (v > 200) & (s < 50)
        
        # 对光晕区域降低亮度
        v_corrected = v.copy().astype(np.float32)
        v_corrected[halo_mask] *= 0.7  # 降低70%的亮度
        v_corrected = np.clip(v_corrected, 0, 255).astype(np.uint8)
        
        # 重新合并
        hsv_corrected = cv2.merge([h, s, v_corrected])
        return cv2.cvtColor(hsv_corrected, cv2.COLOR_HSV2BGR)

    def detect_contours_for_traffic_lights(self, mask, color_name):
        """
        使用轮廓检测方法检测红绿灯
        返回格式: [[[x, y, r], ...]] 与HoughCircles保持一致
        """
        print(f"[轮廓检测] 开始检测 {color_name} 轮廓...")
        
        # 查找轮廓
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        if not contours:
            print(f"[轮廓检测] 未找到 {color_name} 轮廓")
            return None
        
        print(f"[轮廓检测] 找到 {len(contours)} 个 {color_name} 轮廓")
        
        valid_circles = []
        
        for i, contour in enumerate(contours):
            # 计算轮廓面积
            area = cv2.contourArea(contour)
            
            # 计算轮廓的外接圆
            (x, y), radius = cv2.minEnclosingCircle(contour)
            x, y, radius = int(x), int(y), int(radius)
            
            # 计算轮廓的圆形度（4π×面积/周长²）
            perimeter = cv2.arcLength(contour, True)
            if perimeter == 0:
                continue
            
            circularity = 4 * np.pi * area / (perimeter * perimeter)
            
            # 计算轮廓的紧凑度（面积与外接圆面积的比值）
            enclosing_area = np.pi * radius * radius
            compactness = area / enclosing_area if enclosing_area > 0 else 0
            
            print(f"[轮廓检测] 轮廓{i+1}: 中心({x}, {y}), 半径{radius}, 面积{area:.1f}")
            print(f"           圆形度: {circularity:.3f}, 紧凑度: {compactness:.3f}")
            
            valid_circles.append({
                'center': (x, y),
                'radius': radius,
                'area': area,
                'circularity': circularity,
                'compactness': compactness,
                'contour': contour
            })
        
        if not valid_circles:
            print(f"[轮廓检测] 未找到符合条件的 {color_name} 圆形")
            return None
        
        # 转换为与HoughCircles相同的格式: [[[x, y, r], ...]]
        result_circles = []
        for circle in valid_circles:
            x, y = circle['center']
            r = circle['radius']
            result_circles.append([x, y, r])
        
        return [result_circles]  # 返回与HoughCircles相同的格式
    
    def calculate_contour_confidence(self, mask, contours):
        """
        基于轮廓特征计算置信度
        """
        if not contours or len(contours[0]) == 0:
            return 0
        
        confidences = []
        
        for circle_data in contours[0]:
            x, y, r = circle_data
            
            # 创建圆形ROI
            circle_mask = np.zeros(mask.shape, dtype=np.uint8)
            cv2.circle(circle_mask, (x, y), r, 255, -1)
            
            # 计算圆形区域内的像素密度
            roi_mask = cv2.bitwise_and(mask, circle_mask)
            pixel_count = cv2.countNonZero(roi_mask)
            total_pixels = np.pi * r * r
            
            if total_pixels > 0:
                density = pixel_count / total_pixels
                
                # 考虑圆形大小的加权（中等大小的圆形得分更高）
                ideal_radius = (self.min_circle_radius + self.max_circle_radius) / 2
                size_factor = 1.0 - abs(r - ideal_radius) / ideal_radius
                size_factor = max(0.5, size_factor)  # 最小为0.5
                
                # 综合计算置信度
                base_confidence = min(100, density * 120)  # 基础置信度
                weighted_confidence = base_confidence * size_factor
                
                confidences.append(weighted_confidence)
                
                print(f"[置信度计算] 圆形({x}, {y}, {r}): 密度{density:.3f}, 尺寸因子{size_factor:.3f}, 置信度{weighted_confidence:.1f}%")
        
        return max(confidences) if confidences else 0

    def determine_traffic_light_state(self, detection_results):
        """
        根据检测结果判断红绿灯状态
        逻辑：
        1. 如果只检测到红灯轮廓，返回'red'
        2. 如果只检测到绿灯轮廓，返回'green'  
        3. 如果同时检测到红灯和绿灯轮廓，比较面积，返回面积更大的
        4. 如果都没检测到，返回None
        """
        red_area = detection_results['red']['total_area']
        green_area = detection_results['green']['total_area']
        
        has_red = detection_results['red']['contours'] is not None and red_area > 0
        has_green = detection_results['green']['contours'] is not None and green_area > 0
        
        print(f"\n[判断逻辑] 红灯面积: {red_area:.1f}, 绿灯面积: {green_area:.1f}")
        print(f"[判断逻辑] 检测到红灯: {has_red}, 检测到绿灯: {has_green}")
        
        if has_red and not has_green:
            print("[判断逻辑] 只检测到红灯 -> 返回red")
            return 'red'
        elif has_green and not has_red:
            print("[判断逻辑] 只检测到绿灯 -> 返回green")
            return 'green'
        elif has_red and has_green:
            if red_area > green_area:
                print(f"[判断逻辑] 同时检测到红绿灯，红灯面积更大({red_area:.1f} > {green_area:.1f}) -> 返回red")
                return 'red'
            else:
                print(f"[判断逻辑] 同时检测到红绿灯，绿灯面积更大({green_area:.1f} >= {red_area:.1f}) -> 返回green")
                return 'green'
        else:
            print("[判断逻辑] 未检测到任何灯 -> 返回None")
            return None
    
    def calculate_total_contour_area(self, contours):
        """
        计算所有轮廓的总面积
        """
        if not contours or len(contours[0]) == 0:
            return 0
        
        total_area = 0
        for circle_data in contours[0]:
            x, y, r = circle_data
            # 计算圆形面积
            area = np.pi * r * r
            total_area += area
        
        return total_area
    
    def draw_contour_detection(self, image, detection_data, color_name):
        """
        在图像上绘制轮廓检测结果
        """
        if not detection_data['contours']:
            return
        
        print(f"[DEBUG] 绘制{color_name}轮廓检测结果...")
        
        # 颜色映射
        color_map = {
            'red': (0, 0, 255),
            'yellow': (0, 255, 255),
            'green': (0, 255, 0)
        }
        
        color = color_map.get(color_name, (255, 255, 255))
        
        # 绘制所有检测到的圆形
        for i, circle_data in enumerate(detection_data['contours'][0]):
            x, y, r = circle_data
            x, y, r = int(x), int(y), int(r)
            
            # 绘制圆形检测框
            cv2.circle(image, (x, y), r, color, 3)
            # 绘制圆心
            cv2.circle(image, (x, y), 3, color, -1)
            
            print(f"[DEBUG] 绘制{color_name}圆形{i+1}: 中心({x}, {y}), 半径{r}")
        
        # 添加标签
        total_area = detection_data['total_area']
        label = f"{color_name}: {total_area:.1f}"
        
        # 选择第一个圆形的位置作为标签位置
        if detection_data['contours'][0]:
            x, y, r = detection_data['contours'][0][0]
            x, y, r = int(x), int(y), int(r)
            
            label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2)[0]
            
            # 确保标签在图像范围内
            img_height, img_width = image.shape[:2]
            label_x = max(5, min(x - label_size[0] // 2, img_width - label_size[0] - 5))
            label_y = max(35, y - r - 15)
            
            # 绘制标签背景
            cv2.rectangle(image, 
                         (label_x - 5, label_y - 28), 
                         (label_x + label_size[0] + 5, label_y + 5), 
                         (0, 0, 0), -1)
            
            # 绘制标签文字
            cv2.putText(image, label, (label_x, label_y - 5), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
            
            print(f"[DEBUG] 添加{color_name}标签: {label}")

    # ...existing code...
    

if __name__ == '__main__':
    detector = TraditionalTrafficLightDetector()
    detector.run()
